package drds.plus.datasource.connection_restrict;

import drds.plus.datasource.api.NotAvailableException;

import java.util.HashMap;
import java.util.List;

/**
 * 应用连接限制的主要逻辑实现。
 */
public class ConnectionRestrictor {

    /**
     * MAP 策略的应用连接限制, 精确的匹配 Key 和连接槽。
     */
    private HashMap<String, ConnectionRestrictSlot> connectionRestrictSlotKeyToConnectionRestrictSlotMap;

    /**
     * HASH 策略的应用连接限制, 用 Hash + 取模的方式匹配 Key 和连接槽。
     */
    private ConnectionRestrictSlot[] hashConnectionRestrictSlots;

    /**
     * 没有定义业务键 (null Key) 的连接限制槽。
     */
    private ConnectionRestrictSlot nullKeyConnectionRestrictSlot;

    /**
     * 初始化应用连接限制的数据结构, 这些数据结构只会被初始化一次。
     */
    public ConnectionRestrictor(List<ConnectionRestrictEntry> connectionRestrictEntryList) {
        for (ConnectionRestrictEntry connectionRestrictEntry : connectionRestrictEntryList) {
            String[] slotKeys = connectionRestrictEntry.getSlotKeys();
            if (slotKeys.length == 1 && ConnectionRestrictEntry.isNullKey(slotKeys[0])) {
                if (nullKeyConnectionRestrictSlot == null) {
                    nullKeyConnectionRestrictSlot = new ConnectionRestrictSlot(connectionRestrictEntry);
                }
            } else if (slotKeys.length == 1 && ConnectionRestrictEntry.isWildcard(slotKeys[0])) {
                int maxHashSlotSize = connectionRestrictEntry.getMaxHashSlotSize();
                if (maxHashSlotSize < 1) {
                    maxHashSlotSize = 1;
                }
                if (maxHashSlotSize > ConnectionRestrictEntry.MAX_HASH_RESTRICT_SLOT) {
                    maxHashSlotSize = ConnectionRestrictEntry.MAX_HASH_RESTRICT_SLOT;
                }
                if (hashConnectionRestrictSlots == null) {
                    // 每个 HASH 分片都用独立的槽
                    hashConnectionRestrictSlots = new ConnectionRestrictSlot[maxHashSlotSize];
                    for (int i = 0; i < maxHashSlotSize; i++) {
                        hashConnectionRestrictSlots[i] = new ConnectionRestrictSlot(connectionRestrictEntry);
                    }
                }
            } else {
                // 注意, 这里多个业务键同时关联到一个槽
                ConnectionRestrictSlot connectionRestrictSlot = new ConnectionRestrictSlot(connectionRestrictEntry);
                if (connectionRestrictSlotKeyToConnectionRestrictSlotMap == null) {
                    connectionRestrictSlotKeyToConnectionRestrictSlotMap = new HashMap<String, ConnectionRestrictSlot>();
                }
                for (String slotKey : slotKeys) {
                    if (ConnectionRestrictEntry.isNullKey(slotKey)) {
                        if (nullKeyConnectionRestrictSlot == null) {
                            nullKeyConnectionRestrictSlot = connectionRestrictSlot;
                        }
                    } else if (!ConnectionRestrictEntry.isWildcard(slotKey)) {
                        if (!connectionRestrictSlotKeyToConnectionRestrictSlotMap.containsKey(slotKey)) {
                            connectionRestrictSlotKeyToConnectionRestrictSlotMap.put(slotKey, connectionRestrictSlot);
                        }
                    }
                }
            }
        }
    }


    /**
     * 应用连接限制的入口函数。
     */
    public ConnectionRestrictSlot tryAcquire(final int timeoutInMillis) {
        final Object connectRestrictKey = ConnectionRestrictHelper.getConnectRestrictKey();
        ConnectionRestrictSlot connectionRestrictSlot = getConnectionRestrictSlot(connectRestrictKey);
        try {
            // 如果没有匹配的槽, 忽略限制
            if (connectionRestrictSlot != null) {
                if (!connectionRestrictSlot.tryAcquire(timeoutInMillis)) {
                    // 阻塞超时
                    throw new NotAvailableException("No connection available for '" + connectRestrictKey + "' within configured blocking timeout (" + timeoutInMillis + "[ms])");
                }
            }
        } catch (InterruptedException e) {
            throw new NotAvailableException("Allocate connection for '" + connectRestrictKey + "' interrupted within configured blocking timeout (" + timeoutInMillis + "[ms]) , caused by " + e.getCause());
        } catch (RuntimeException e) {
            throw new NotAvailableException("Allocate connection for '" + connectRestrictKey + "' failed: unexpected exception ");
        }
        return connectionRestrictSlot;
    }

    /**
     * 从数据结构中查找应用连接限制的槽。
     */
    public ConnectionRestrictSlot getConnectionRestrictSlot(Object connectionRestrictSlotKey) {
        if (connectionRestrictSlotKey != null) {
            ConnectionRestrictSlot connectionRestrictSlot = null;
            if (connectionRestrictSlotKeyToConnectionRestrictSlotMap != null) {
                // 首先精确匹配
                connectionRestrictSlot = connectionRestrictSlotKeyToConnectionRestrictSlotMap.get(String.valueOf(connectionRestrictSlotKey));
            }
            if (connectionRestrictSlot == null) {
                if (hashConnectionRestrictSlots != null) {
                    // 如果没有精确指定, 则用 HASH 方式
                    final int mod = Math.abs(connectionRestrictSlotKey.hashCode() % hashConnectionRestrictSlots.length);
                    connectionRestrictSlot = hashConnectionRestrictSlots[mod];
                }
            }
            return connectionRestrictSlot;
        }
        // 没有定义业务键, 走 null Key 槽
        return nullKeyConnectionRestrictSlot;
    }

}
